import numpy as np
from easy_io import read_pkl_file, write_pkl_file
from merge_candidates import merge
from label_candidates import label, prepare_labelinfo


def index(candidates, start=0):
    for i, c in enumerate(candidates, start=start):
        c.update(index=i)
    return candidates


def path(candidates):
    for c in candidates:
        c.update(path=str(c['index']))
    return candidates


def divide(candidates, divisions):
    for c in candidates:
        c.update(fold=divisions[c['scanid']])
    return candidates


def calculate(candidates):
    for c in candidates:
        bbox = np.asarray(c['bbox'])
        center = (bbox[:3] + bbox[3:] - 1) / 2
        shape = bbox[3:] - bbox[:3]
        diameter = np.max(shape[:2])
        c.update(center=center, diameter=diameter)
    return candidates


def pipeline(candidates, apply_merge=True, **kwargs):
    if apply_merge:
        candidates = merge(candidates)
    if 'pos_labelinfo' in kwargs and kwargs['pos_labelinfo'] is not None:
        candidates = label(candidates, kwargs['pos_labelinfo'], kwargs.get('mid_labelinfo', None))
    candidates = index(candidates, kwargs.get('start', 0))
    candidates = path(candidates)
    if 'divisions' in kwargs and kwargs['divisions'] is not None:
        candidates = divide(candidates, kwargs['divisions'])
    candidates = calculate(candidates)
    return candidates


def main(infile, outfile, apply_merge=True, division_file=None,
         pos_labelinfo_file=None, pos_label_dict_file=None,
         mid_labelinfo_file=None, mid_label_dict_file=None,
         start=0):
    divisions = (division_file and read_pkl_file(division_file)) or None
    pos_labelinfo = (pos_labelinfo_file and read_pkl_file(pos_labelinfo_file)) or \
                    (pos_label_dict_file and prepare_labelinfo(read_pkl_file(pos_label_dict_file))) or None
    mid_labelinfo = (mid_labelinfo_file and read_pkl_file(mid_labelinfo_file)) or \
                    (mid_label_dict_file and prepare_labelinfo(read_pkl_file(mid_label_dict_file))) or None

    candidates = read_pkl_file(infile)
    candidates = pipeline(candidates, apply_merge=apply_merge, divisions=divisions, pos_labelinfo=pos_labelinfo, mid_labelinfo=mid_labelinfo,
                          start=start)
    write_pkl_file(outfile, candidates)


if __name__ == '__main__':
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v6.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v7.pkl',
    #     division_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_divisions.pkl',
    #     pos_label_dict_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_label_dict_3_backup.pkl',
    #     mid_label_dict_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_label_dict_1_backup.pkl',
    # )
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/spie_candidates_v1.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/spie_candidates_v2.pkl',
    #     start=30000,
    # )
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v1.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v2.pkl',
    #     start=60000,
    # )
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v1.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v2.pkl',
    #     start=90000,
    # )

    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v8.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v9.pkl',
    #     division_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_divisions.pkl',
    #     pos_label_dict_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_label_dict_3_backup.pkl',
    #     mid_label_dict_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_label_dict_1_backup.pkl',
    # )
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/spie_candidates_v3.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/spie_candidates_v4.pkl',
    #     start=30000,
    # )
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v3.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v4.pkl',
    #     start=60000,
    # )
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v3.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v4.pkl',
    #     start=90000,
    # )

    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v8.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v9_nomerge.pkl',
    #     apply_merge=False,
    #     division_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_divisions.pkl',
    #     pos_label_dict_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_label_dict_3_backup.pkl',
    #     mid_label_dict_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_label_dict_1_backup.pkl',
    # )

    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v10.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v11.pkl',
    #     apply_merge=True,
    #     division_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_divisions.pkl',
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_3.pkl',
    #     mid_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_1.pkl',
    # )
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/spie_candidates_v5.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/spie_candidates_v6_nomerge.pkl',
    #     apply_merge=False,
    #     start=120000,
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/spie_labelinfo.pkl'
    # )
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v5.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v6_nomerge.pkl',
    #     apply_merge=False,
    #     start=60000,
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_3.pkl',
    # )
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v5.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v6_nomerge.pkl',
    #     apply_merge=False,
    #     start=90000,
    # )

    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v12.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v13.pkl',
    #     apply_merge=True,
    #     division_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_divisions.pkl',
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_3.pkl',
    #     mid_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_1.pkl',
    # )

    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v12.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v13_nomerge.pkl',
    #     apply_merge=False,
    #     division_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_divisions.pkl',
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_3.pkl',
    #     mid_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_1.pkl',
    # )

    # divisions = read_pkl_file('/ssd_1t/huzq/kaggle_data/lidc_kaggle_divisions_4fold.pkl')
    # write_pkl_file('/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v13.pkl',
    #                divide(
    #                    read_pkl_file('/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v13.pkl'),
    #                    divisions
    #                ))

    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/spie_candidates_v7.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/spie_candidates_v8_nomerge.pkl',
    #     apply_merge=False,
    #     start=300000,
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/spie_labelinfo.pkl'
    # )
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v7.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v8_nomerge.pkl',
    #     apply_merge=False,
    #     start=200000,
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_3.pkl',
    # )
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v7.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v8_nomerge.pkl',
    #     apply_merge=False,
    #     start=100000,
    # )

    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v14.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v15.pkl',
    #     apply_merge=False,
    #     division_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_divisions_4fold.pkl',
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_3.pkl',
    #     mid_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_1.pkl',
    # )

    # candidates = read_pkl_file('/ssd_1t/huzq/kaggle_data/spie_candidates_v8_nomerge.pkl')
    # for c in candidates:
    #     c.update(source='spie')
    # write_pkl_file('/ssd_1t/huzq/kaggle_data/spie_candidates_v8_nomerge.pkl', candidates)

    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/spie_candidates_v9.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/spie_candidates_v10.pkl',
    #     apply_merge=False,
    #     start=300000,
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/spie_labelinfo.pkl'
    # )
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v9.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v10.pkl',
    #     apply_merge=False,
    #     start=200000,
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_3.pkl',
    # )
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v9.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v10.pkl',
    #     apply_merge=False,
    #     start=100000,
    # )

    candidates = read_pkl_file('/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v10.pkl')
    for c in candidates:
        c.update(source='kaggle_testset')
    write_pkl_file('/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v10.pkl', candidates)
    candidates = read_pkl_file('/ssd_1t/huzq/kaggle_data/spie_candidates_v10.pkl')
    for c in candidates:
        c.update(source='spie')
    write_pkl_file('/ssd_1t/huzq/kaggle_data/spie_candidates_v10.pkl', candidates)
